"""
Created on 2017 10.17
@author: liupeng
"""

import numpy as np
import tensorflow as tf
from tensorflow.contrib.slim.python.slim.nets.resnet_utils import resnet_arg_scope
from tensorflow.contrib.slim.python.slim.nets.resnet_v1 import resnet_v1_101

slim = tf.contrib.slim
import numpy as np
import argparse
import os
from PIL import Image
from datetime import datetime
import math
import time

export_path = "/home/dolly/Desktop/resnet101"

batch_size = 32
height, width = 224, 224
X = tf.placeholder(tf.float32, [None, height, width, 3], name="shuru")
Y = tf.placeholder(tf.float32, [None, 10], name="shuchu")

builder = tf.saved_model.builder.SavedModelBuilder(export_path)
tensor_info_x = tf.saved_model.utils.build_tensor_info(X)
tensor_info_y = tf.saved_model.utils.build_tensor_info(Y)

prediction_signature = (
    tf.saved_model.signature_def_utils.build_signature_def(
        inputs={'input': tensor_info_x},
        outputs={'output': tensor_info_y},
        method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME))

legacy_init_op = tf.group(tf.tables_initializer(), name='legacy_init_op')

print("-----------------------------main.py start--------------------------")

def test():
    # model
    arg_scope = resnet_arg_scope()
    with slim.arg_scope(arg_scope):
        net, end_points = resnet_v1_101(X, is_training=False)

    sess = tf.Session()

    # reload model
    saver1 = tf.train.Saver(tf.global_variables())
    checkpoint_path = '/home/dolly/checkpoints/resnet_v1_101_2016_08_28/resnet_v1_101.ckpt'
    saver1.restore(sess, checkpoint_path)

    num_classes = 10
    net = slim.conv2d(net, num_classes, [1, 1], activation_fn=None, normalizer_fn=None, scope='logits2')
    net = tf.squeeze(net, [1, 2], name='SpatialSqueeze')
    # initializer
    init = tf.global_variables_initializer()
    sess.run(init)

    # saver2 = tf.train.Saver(tf.global_variables())
    # saver2.restore(sess, "model/101/fine-tune-1000")

    # input
    # input = X
    # inputs = tf.random_uniform((batch_size, height, width, 3))

    # 对图片输入进行预处理
    im = tf.read_file("/home/dolly/Desktop/img1.jpg")
    im = tf.image.decode_jpeg(im)
    im = tf.image.resize_images(im, (width, height))
    im = tf.reshape(im, [-1, height, width, 3])
    im = tf.cast(im, tf.float32)
    inputs = im

    # run
    # images = sess.run(inputs)
    # print(type(images))
    # print(images)
    # start_time = time.time()
    # out_put = sess.run(net, feed_dict={X: images})
    # print(out_put)

    # 计算推理时间
    # duration = time.time() - start_time

    # reshape输出
    # predict = tf.reshape(out_put, [-1, num_classes])
    # max_idx_p = tf.argmax(predict, 1)
    # print(out_put.shape)
    # print(sess.run(max_idx_p))
    # print('run time:', duration)

    # 导出
    builder.add_meta_graph_and_variables(
        sess, [tf.saved_model.tag_constants.SERVING],
        signature_def_map={
            # 'prediction':
            #     prediction_signature,
            tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
                prediction_signature,
        },
        legacy_init_op=legacy_init_op)
    builder.save()
    sess.close()


test()
